# rm(list=ls());
gc()

# library(dplyr)
# library(tidyr)
# library(loo)

out.fldr <- "results/us/"

id_info_full <- data.frame()
id_est_full <- weights_est_full <- gamma_est_full <- ns_est_full <- cor_est_full <- sal_est_full <- old_est_full <- loo_est_full <- prcomp_est_full <- data.frame()

for (i in 79:115){
  load(paste0(out.fldr, i, "/sirt.Rda"))
  fit_sep <- fit_sum
  loo_sep <- loo_sum
  rm(fit_sum)
  rm(loo_sum)
  load(paste0(out.fldr, i, "/nsirt.Rda"))
  fit_nsep <- fit_sum
  loo_nsep <- loo_sum
  rm(fit_sum)
  rm(loo_sum)
  
  id_info_full <- bind_rows(id_info_full, mem_desc %>% 
                              arrange(id) %>% 
                              select(id, icpsr, bioname, party_code, state_abbrev, party, confed) %>% mutate(session = i))
  
  id_est <- bind_rows(fit_nsep %>% 
                        filter(parameter == "theta") %>% 
                        mutate(model = "nonsep",
                               id = rep(mem_desc$icpsr, 2),
                               dim = c(rep(1, nrow(mem_desc)), rep(2, nrow(mem_desc))),
                               member = c(mem_desc$party, mem_desc$confed)),
                      fit_sep %>% 
                        filter(parameter == "theta") %>% 
                        mutate(model = "sep",
                               id = rep(mem_desc$icpsr, 2),
                               dim = c(rep(1, nrow(mem_desc)), rep(2, nrow(mem_desc))),
                               member = c(mem_desc$id, mem_desc$confed))
  )
  id_est_full <- bind_rows(id_est_full, id_est %>% mutate(session = i))
  
  ns_est <- fit_nsep %>% 
    filter(parameter == "DL") %>% 
    slice(2)
  ns_est_full <- bind_rows(ns_est_full, ns_est %>% mutate(session = i))
  
  weights_est <- bind_rows(fit_nsep %>% 
                             filter(parameter == "weights") %>% 
                             mutate(model = "nonsep", id = 1:4),
                           fit_sep %>% 
                             filter(parameter == "weights") %>% 
                             mutate(model = "sep", id = 1:4))
  weights_est_full <- bind_rows(weights_est_full, weights_est %>% mutate(session = i))
  
  sal_est <- bind_rows(fit_nsep %>% 
                         filter(parameter == "sal") %>% 
                         mutate(model = "nonsep", id = 1:2),
                       fit_sep %>% 
                         filter(parameter == "sal") %>% 
                         mutate(model = "sep", id = 1:2))
  sal_est_full <- bind_rows(sal_est_full, sal_est %>% mutate(session = i))
  
  old_est <- bind_rows(fit_nsep %>% 
                         filter(parameter == "r_old") %>% 
                         mutate(model = "nonsep", id = 1:2),
                       fit_sep %>% 
                         filter(parameter == "r_old") %>% 
                         mutate(model = "sep", id = 1:2))
  old_est_full <- bind_rows(old_est_full, old_est %>% mutate(session = i))
  
  cor_est <- bind_rows(fit_nsep %>% 
                         filter(parameter == "r_dim") %>% 
                         mutate(model = "nonsep"),
                       fit_sep %>% 
                         filter(parameter == "r_dim") %>% 
                         mutate(model = "sep"))
  cor_est_full <- bind_rows(cor_est_full, cor_est %>% mutate(session = i))
  
  prcomp_est <- bind_rows(fit_nsep %>% 
                            filter(parameter %in% c("p_share", "m_prop_compl", "sd_prop_compl")) %>% 
                            mutate(model = "nonsep"),
                          fit_sep %>% 
                            filter(parameter %in% c("p_share", "m_prop_compl", "sd_prop_compl")) %>% 
                            mutate(model = "sep"))
  prcomp_est_full <- bind_rows(prcomp_est_full, prcomp_est %>% mutate(session = i))
  
  gamma_est <- bind_rows(fit_nsep %>% 
                           filter(parameter == "gamma_post") %>% 
                           mutate(dim = ifelse(i == 79, 1:2, 0:2), model = "nonsep"),
                         fit_sep %>% 
                           filter(parameter == "gamma_post") %>% 
                           mutate(dim = ifelse(i == 79, 1:2, 0:2), model = "sep"),
                         fit_nsep %>% 
                           filter(parameter == "r_sq") %>% 
                           mutate(dim = c(1,2), model = "nonsep"),
                         fit_sep %>% 
                           filter(parameter == "r_sq") %>% 
                           mutate(dim = c(1,2), model = "sep"))
  gamma_est_full <- bind_rows(gamma_est_full, gamma_est %>% mutate(session = i))
  
  loo_comp <- data.frame(loo_compare(loo_sep, loo_nsep)) %>% 
    mutate(model = rownames(.),
           model = case_when(model == "model1" ~ "sep",
                             TRUE ~ "nonsep"),
           session = i)
  loo_est_full <- bind_rows(loo_est_full, loo_comp)
}

save(id_info_full,
     id_est_full, weights_est_full, gamma_est_full, ns_est_full, sal_est_full, old_est_full, cor_est_full, loo_est_full, prcomp_est_full, 
     file = paste0(out.fldr, "/full.Rda"))
